from util.embeddings import embedding_filter
from util.entity import Article
from util.es import query_data


def es_recall(queries: list[str]):
    ds = query_data(index_name="law1", search_terms=queries, fields=["title", "content"], size=100)
    return ds


def file_filter(recall_result, filter_files):
    if filter_files is None:
        return recall_result
    if not filter_files or len(filter_files) == 0:
        return []
    result = []
    for r in recall_result:
        if r["file_id"] in filter_files:
            result.append(r)

    return result


def permission_filter(recall_result, filter_permissions):
    if not filter_permissions or len(filter_permissions) == 0:
        return []
    result = []
    for r in recall_result:
        if bool(set(r["permission_set"]) & set(filter_permissions)):
            result.append(r)
    return result


def result_filter(queries, recall_result, filter_files, permission_list):
    if filter_files is not None:
        recall_result = file_filter(recall_result, filter_files)
    if permission_list is not None:
        recall_result = permission_filter(recall_result, permission_list)
    embed_filter_result = embedding_filter(queries, recall_result, threshold=0.4)
    return embed_filter_result


def recall(queries: list[str], filter_files, permission_list, q_type) -> list[dict]:
    es_result = es_recall(queries)
    es_result = result_filter(queries, es_result, filter_files, permission_list)
    es_result = [
        {"recall_type": "embedding",
         "result": Article(name=x["title"], content=x["content"], score=x["score"], file_id=x["file_id"],
                           permission_type=x.get("permission_type", 1), permission_set=x.get("permission_set", []))}
        for x in es_result]
    print("recall result: {}".format(es_result))
    return es_result
